from airflow import DAG
# 将第2行的BashOperator导入替换为
from airflow.providers.standard.operators.bash import BashOperator

# 将PythonOperator导入替换为
from airflow.providers.standard.operators.python import PythonOperator
from datetime import datetime, timedelta
import os
import shutil
import time
import requests

# 定义默认参数
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2023, 10, 1),
    'email': ['airflow@example.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

# 定义DAG
dag = DAG(
    dag_id='daily_model_update',
    schedule='0 0 * * *',  # 每天凌晨0点执行
    start_date=datetime(2023, 1, 1),
    catchup=False,
    default_args=default_args
)


# 任务1: 生成用户特征
generate_user_features_task = BashOperator(
    task_id='generate_user_features',
    bash_command='''
    python /code/guocu/feature/user_inter_click.py
    python /code/guocu/feature/user_inter_login.py
    python /code/guocu/feature/user_feature_merge.py
    ''',
    dag=dag,
)

# 任务2: 生成物品特征
generate_item_features_task = BashOperator(
    task_id='generate_item_features',
    bash_command='''
    python /code/guocu/feature/item_inter_num.py
    python /code/guocu/feature/user_item_inter_rank.py
    python /code/guocu/feature/iterm_feature_merge.py
    ''',
    dag=dag,
)

# 任务3: 预处理全部数据
preprocess_data_task = BashOperator(
    task_id='preprocess_all_data',
    bash_command='''
    python /code/guocu/data_preprocessing/click2exposure.py
    python /code/guocu/data_preprocessing/DataProcessing.py
    ''',
    dag=dag,
)

# 任务4: 模型训练及保存
def train_model():
    # 备份旧模型
    old_model_path = '/data/GuoCu_data/models/final_lgb_model.txt'
    if os.path.exists(old_model_path):
        backup_path = f'{old_model_path}.{datetime.now().strftime("%Y%m%d%H%M%S")}.bak'
        shutil.copy2(old_model_path, backup_path)
        print(f'已备份旧模型至: {backup_path}')
    
    # 调用训练脚本
    os.system('python /code/guocu/train/train_lightgbm.py')


train_model_task = PythonOperator(
    task_id='train_and_save_model',
    python_callable=train_model,
    dag=dag,
)

# 任务6: 部署的模型更新
def update_deployed_model(timeout=300, check_interval=5):
    """
    更新部署的模型
    :param timeout: 最大等待时间(秒)，默认为300秒(5分钟)
    :param check_interval: 检查间隔(秒)，默认为5秒
    :return: tuple (是否成功, 消息)
    """
    # 调用API来通知服务更新模型
    try:
        print("开始触发模型更新...")
        response = requests.post('http://localhost:8000/reload-model')
        if response.status_code != 200:
            msg = f'触发模型更新失败: {response.status_code}, {response.text}'
            print(msg)
            return False, msg
        
        print("模型更新已触发，等待加载完成...")
        start_time = time.time()
        
        # 轮询检查模型加载状态
        while time.time() - start_time < timeout:
            time.sleep(check_interval)
            try:
                status_response = requests.get('http://localhost:8000/model-status')
                if status_response.status_code == 200:
                    status_data = status_response.json()
                    if status_data['status'] == 'ready':
                        msg = '模型已成功更新并加载完成'
                        print(msg)
                        return True, msg
                    elif status_data['status'] == 'loading':
                        print(f"模型加载中: {status_data['message']}")
                    else:
                        msg = f'未知状态: {status_data}'
                        print(msg)
                else:
                    msg = f'检查模型状态失败: {status_response.status_code}'
                    print(msg)
            except Exception as e:
                msg = f'检查模型状态异常: {str(e)}'
                print(msg)
        
        # 超时
        msg = f'模型更新超时({timeout}秒)'
        print(msg)
        return False, msg
    
    except Exception as e:
        msg = f'模型更新异常: {str(e)}'
        print(msg)
        return False, msg


update_deployed_model_task = PythonOperator(
    task_id='update_deployed_model',
    python_callable=update_deployed_model,
    dag=dag,
)

# 设置任务依赖关系
[generate_user_features_task, generate_item_features_task] >> preprocess_data_task
preprocess_data_task >> train_model_task
train_model_task >> update_deployed_model_task